import pymol
from pymol import cmd
#import argparse
import sys

''' parser = argparse.ArgumentParser(description='draw pdb image')
parser.add_argument('-ir', type=str, dest='raw_pdb_file', help='raw PDB file')
parser.add_argument('-im', type=str, dest='mu_pdb_file', help='mutated PDB file')
parser.add_argument('-o', default='result.png', dest='output_img', type=str, help='the result PDB image')
parser.add_argument('-c', default='A', dest='chain', type=str, help='chain which the mutation located')
parser.add_argument('-r', dest='res_no', type=str, help='residue number') '''


class PDBViewer:
    name_list = None

    def __init__(self):
        pymol.pymol_argv = ['pymol', '-qc']
        pymol.finish_launching()

    def load(self, pdb_file, name=''):
        if name:
            cmd.load(pdb_file, object=name)
        else:
            cmd.load(pdb_file)

    def load_pdbs(self, pdb_file_list, name_li=None):
        if not name_li or not len(name_li) == len(pdb_file_list):
            name_li = ['pdb-' + str(x) for x in range(len(pdb_file_list))]
        for pdb, name in zip(pdb_file_list, name_li):
            #print pdb, name
            cmd.load(pdb, object=name)
        self.name_list = name_li

    def set_style(self, stick_models):
        cmd.hide('all')
        cmd.show('cartoon')
        cmd.spectrum()
        for el in stick_models:
            cmd.show('sticks', el)
        #cmd.zoom(stick_models[0], 20)

    def center(self, model):
        cmd.center(model)

    def make_model(self, model_name, selection):
        #print selection
        cmd.select(model_name, selection)

    def save_img(self, filename, width=800, height=800):
        # cmd.save(filename=filename, format='png')
        # cmd.ray(width=800, height=800)
        cmd.select('none')
        cmd.png(filename=filename, width=width, height=height, ray=1)



#
# pv = PDBViewer()
# pdb_file = '/home/zhonghua/test/pdbs/1a1e.pdb'
#
# pv.load(pdb_file=pdb_file)
# pv.make_model('test', '///A/159')
# pv.set_style(['test'])
# pv.center('test')
# pv.save_img('/home/zhonghua/test.png', 1000, 1000)
print sys.argv
raw_pdb, mutant_pdb, chain, res_no, output_img = sys.argv[1:]

pv = PDBViewer()
pv.load_pdbs([raw_pdb, mutant_pdb])
res_list = []
for el in pv.name_list:
    res_list.append(el+'res')
    pv.make_model(el+'res', '/{0}//{1}/{2}'.format(el, chain, res_no))
pv.set_style(res_list)
pv.save_img(output_img)
